# 导入tushare
import tushare as ts

# 初始化pro接口
pro = ts.pro_api('31d3906685a1b0d31d52b4478cc1448a5ee4235abe9f55ab0a8943b9')

# 拉取数据
df = pro.daily(**{
    "ts_code": "",
    "trade_date": "",
    "start_date": 202500321,
    "end_date": 20250331,
    "offset": "",
    "limit": ""
}, fields=[
    "ts_code",
    "trade_date",
    "open",
    "high",
    "low",
    "close",
    "pre_close",
    "change",
    "pct_chg",
    "vol",
    "amount"
])
print(df)

# 数据统计
stats = df.describe()
print(stats)
# 按股票代码分组计算
df_grouped = df.groupby('ts_code')

# 计算MA5和MA10
df['ma5'] = df_grouped['close'].transform(lambda x: x.rolling(5).mean())
df['ma10'] = df_grouped['close'].transform(lambda x: x.rolling(10).mean())


# 计算RSI（以14日为例）
def calculate_rsi(data, window=14):
    delta = data.diff()
    gain = (delta.where(delta > 0, 0)).rolling(window).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window).mean()
    rs = gain / loss
    return 100 - (100 / (1 + rs))


df['rsi'] = df_grouped['close'].transform(calculate_rsi)
# 确保数据按日期排序
df = df.sort_values(by='trade_date')

# 判断上穿条件
df['ma5_prev'] = df_grouped['ma5'].shift(1)
df['ma10_prev'] = df_grouped['ma10'].shift(1)
buy_signals = df[(df['ma5'] > df['ma10']) & (df['ma5_prev'] <= df['ma10_prev'])]

# 保存结果
buy_signals.to_csv('buy_signals.csv')
